from typing import Annotated

import os

from sqlmodel import create_engine, SQLModel, Session, select
from sqlmodel.pool import StaticPool
from fastapi import Depends

from . import models  # noqa: F401
from .settings import settings

DATABASE_URL = os.getenv("DATABASE_URL") or settings.DATABASE_URL
# 确保 DATABASE_URL 不为空，否则抛出异常
if DATABASE_URL is None:
    raise ValueError("DATABASE_URL must not be None")


# 初始化数据库引擎
def get_engine(database_url: str = DATABASE_URL):
    if database_url.startswith("sqlite:///:memory:"):
        return create_engine(
            database_url,
            connect_args={"check_same_thread": False},
            poolclass=StaticPool,
        )
    return create_engine(database_url)


engine = get_engine(DATABASE_URL)


def create_db_and_tables():
    SQLModel.metadata.create_all(engine)


def get_session():
    with Session(engine) as session:
        yield session


SessionDep = Annotated[Session, Depends(get_session)]


def exists(session: Session, model: type[SQLModel], **kwargs):
    return session.exec(select(model).filter_by(**kwargs)).first() is not None
